import paddle
import paddle.nn.functional as F
import paddle.static as static
import paddleslim as slim
import numpy as np
import paddle.vision.transforms as T

# ----------- 1. 导入依赖。已在上面导入 ----------- #

# ----------- 2. 初始化 SANAS 搜索实例 ----------- #
def init_sanas():
    port = np.random.randint(8337, 8773)
    return slim.nas.SANAS(configs=[('MobileNetV2Space')], server_addr=("", port), save_checkpoint=None)

# ----------- 3. 构建网络程序 ----------- #
def build_program(archs):
    paddle.enable_static()
    train_program = static.Program()
    startup_program = static.Program()
    with static.program_guard(train_program, startup_program):
        data = static.data(name='data', shape=[None, 3, 32, 32], dtype='float32')
        label = static.data(name='label', shape=[None, 1], dtype='int64')
        gt = paddle.reshape(label, [-1, 1])
        output = archs(data)
        output = static.nn.fc(output, size=10)
        softmax_out = F.softmax(output)
        cost = F.cross_entropy(softmax_out, label=gt)
        avg_cost = paddle.mean(cost)
        acc_top1 = paddle.metric.accuracy(input=softmax_out, label=gt, k=1)
        acc_top5 = paddle.metric.accuracy(input=softmax_out, label=gt, k=5)
        test_program = static.default_main_program().clone(for_test=True)
        optimizer = paddle.optimizer.Adam(learning_rate=0.1)
        optimizer.minimize(avg_cost)
        place = paddle.CUDAPlace(0)
        exe = static.Executor(place)
        exe.run(startup_program)
    return exe, train_program, test_program, (data, label), avg_cost, acc_top1, acc_top5

# ----------- 4. 定义数据输入 ----------- #
def input_data(image, label):
    transform = T.Compose([T.Transpose(), T.Normalize([127.5], [127.5])])
    train_dataset = paddle.vision.datasets.Cifar10(mode="train", transform=transform, backend='cv2')
    train_loader = paddle.io.DataLoader(
        train_dataset,
        places=paddle.CUDAPlace(0),
        feed_list=[image, label],
        drop_last=True,
        batch_size=4096,
        return_list=False,
        shuffle=True
    )
    eval_dataset = paddle.vision.datasets.Cifar10(mode="test", transform=transform, backend='cv2')
    eval_loader = paddle.io.DataLoader(
        eval_dataset,
        places=paddle.CPUPlace(),
        feed_list=[image, label],
        drop_last=False,
        batch_size=512,
        return_list=False,
        shuffle=False
    )
    return train_loader, eval_loader

# ----------- 5. 训练函数 ----------- #
def start_train(exe, program, data_loader, avg_cost, acc_top1, acc_top5):
    outputs = [avg_cost.name, acc_top1.name, acc_top5.name]
    for data in data_loader():
        batch_reward = exe.run(program, feed=data, fetch_list=outputs)
        print("TRAIN: loss: {}, acc1: {}, acc5:{}".format(batch_reward[0], batch_reward[1], batch_reward[2]))

# ----------- 6. 评估函数 ----------- #
def start_eval(exe, program, data_loader, avg_cost, acc_top1, acc_top5):
    losses, acc1s, acc5s = [], [], []
    outputs = [avg_cost.name, acc_top1.name, acc_top5.name]
    for data in data_loader():
        batch_loss, batch_acc1, batch_acc5 = exe.run(program, feed=data, fetch_list=outputs)
        losses.append(batch_loss)
        acc1s.append(batch_acc1)
        acc5s.append(batch_acc5)
        print("TEST: loss: {}, acc1: {}, acc5:{}".format(batch_loss, batch_acc1, batch_acc5))
    # 分别统计均值
    avg_cost_val = np.mean(losses)
    avg_acc1_val = np.mean(acc1s)
    avg_acc5_val = np.mean(acc5s)
    print("FINAL TEST: avg_cost: {}, acc1: {}, acc5: {}".format(avg_cost_val, avg_acc1_val, avg_acc5_val))
    return [avg_cost_val, avg_acc1_val, avg_acc5_val]

# ----------- 7. 主流程整合 ----------- #
def main():
    sanas = init_sanas()
    max_steps = 10        # 搜索步数
    epochs_per_model = 30 # 每次架构训练epoch数量
    FLOPS_LIMIT = 6555276

    for step in range(max_steps):
        archs = sanas.next_archs()[0]
        exe, train_program, eval_program, (images, label), avg_cost, acc_top1, acc_top5 = build_program(archs)
        train_loader, eval_loader = input_data(images, label)

        current_flops = slim.analysis.flops(train_program)
        if current_flops > FLOPS_LIMIT:
            print(f"Skip arch with FLOPs {current_flops}")
            continue

        for epoch in range(epochs_per_model):
            print(f"Step {step+1}/{max_steps}, Epoch {epoch+1}/{epochs_per_model}")
            start_train(exe, train_program, train_loader, avg_cost, acc_top1, acc_top5)

        finally_reward = start_eval(exe, eval_program, eval_loader, avg_cost, acc_top1, acc_top5)
        sanas.reward(float(finally_reward[1]))

if __name__ == "__main__":
    main()